{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f19fe218-8272-4a78-95dc-b45c7944d26d",
   "metadata": {},
   "source": [
    "# Building and deploying machine learning solutions with Vertex AI: Challenge Lab"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d91e5394-d8e8-4b56-99a1-f7c3b0f574f4",
   "metadata": {},
   "source": [
    "This Challenge Lab is recommended for students who have enrolled in **Building and deploying machine learning solutions with Vertex AI**. You will be given a scenario and a set of tasks. Instead of following step-by-step instructions, you will use the skills learned from the labs in the quest to figure out how to complete the tasks on your own! An automated scoring system (shown on the lab instructions page) will provide feedback on whether you have completed your tasks correctly.\n",
    "\n",
    "When you take a Challenge Lab, you will not be taught Google Cloud concepts. To build the solution to the challenge presented, use skills learned from the labs in the Quest this challenge lab is part of. You are expected to extend your learned skills and complete all the **`TODO:`** comments in this notebook.\n",
    "\n",
    "Are you ready for the challenge?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4908fb9b-2048-48fc-a42c-2fdf76aea51e",
   "metadata": {},
   "source": [
    "## Scenario"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eefbdce5-4287-4740-bdbd-729d15d8ab7f",
   "metadata": {},
   "source": [
    "You were recently hired as a Machine Learning Engineer at a startup movie review website. Your manager has tasked you with building a machine learning model to classify the sentiment of user movie reviews as positive or negative. These predictions will be used as an input in downstream movie rating systems and to surface top supportive and critical reviews on the movie website application. The challenge: your business requirements are that you have just 6 weeks to productionize a model that achieves great than 75% accuracy to improve upon an existing bootstrapped solution. Furthermore, after doing some exploratory analysis in your startup's data warehouse, you found that you only have a small dataset of 50k text reviews to build a higher performing solution.\n",
    "\n",
    "To build and deploy a high performance machine learning model with limited data quickly, you will walk through training and deploying a custom TensorFlow BERT sentiment classifier for online predictions on Google Cloud's [Vertex AI](https://cloud.google.com/vertex-ai) platform. Vertex AI is Google Cloud's next generation machine learning development platform where you can leverage the latest ML pre-built components and AutoML to significantly enhance your development productivity, scale your workflow and decision making with your data, and accelerate time to value.\n",
    "\n",
    "![Vertex AI: Challenge Lab](./images/vertex-challenge-lab.png \"Vertex Challenge Lab\")\n",
    "\n",
    "First, you will progress through a typical experimentation workflow where you will build your model from pre-trained BERT components from TF-Hub and `tf.keras` classification layers to train and evaluate your model in a Vertex Notebook. You will then package your model code into a Docker container to train on Google Cloud's Vertex AI. Lastly, you will define and run a Kubeflow Pipeline on Vertex Pipelines that trains and deploys your model to a Vertex Endpoint that you will query for online predictions."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8955d75d-cfa4-43af-8783-d2aec5ae525e",
   "metadata": {},
   "source": [
    "## Learning objectives"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b386b37c-2ce1-4b1f-8c90-b83bda6075c8",
   "metadata": {},
   "source": [
    "* Train a TensorFlow model locally in a hosted [**Vertex Notebook**](https://cloud.google.com/vertex-ai/docs/general/notebooks?hl=sv).\n",
    "* Containerize your training code with [**Cloud Build**](https://cloud.google.com/build) and push it to [**Google Cloud Artifact Registry**](https://cloud.google.com/artifact-registry).\n",
    "* Define a pipeline using the [**Kubeflow Pipelines (KFP) V2 SDK**](https://www.kubeflow.org/docs/components/pipelines/sdk/v2/v2-compatibility) to train and deploy your model on [**Vertex Pipelines**](https://cloud.google.com/vertex-ai/docs/pipelines).\n",
    "* Query your model on a [**Vertex Endpoint**](https://cloud.google.com/vertex-ai/docs/predictions/getting-predictions) using online predictions."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d23538a-e809-4747-9bd4-5610f8544ea1",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d23538a-e809-4747-9bd4-5610f8544ea1",
   "metadata": {},
   "source": [
    "**NOTE: Make sure you have installed the required packages for the lab as specified in the Task 2 > step 3 of the lab instructions.**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23626dae",
   "metadata": {},
   "source": [
    "### Install libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66c560eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip3 install google-cloud-bigquery==3.25.0 -U\n",
    "!pip install google-cloud-aiplatform==1.59.0\n",
    "!pip uninstall -y shapely pygeos geopandas\n",
    "!pip install shapely==1.8.5.post1 pygeos==0.12.0 geopandas>=0.12.2\n",
    "# Install pydot and graphviz\n",
    "!pip install pydot\n",
    "!sudo apt install graphviz -y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "591e3c3e",
   "metadata": {},
   "source": [
    "### After installing these packages you'll need to restart the kernel:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb03837e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "if not os.getenv(\"IS_TESTING\"):\n",
    "    # Automatically restart kernel after installs\n",
    "    import IPython\n",
    "\n",
    "    app = IPython.Application.instance()\n",
    "    app.kernel.do_shutdown(True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c4480c8-710c-40dd-93c2-c51e67e59760",
   "metadata": {},
   "source": [
    "### Define constants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e0feaf4-9849-4636-b736-d3cd8a051579",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add installed library dependencies to Python PATH variable.\n",
    "PATH=%env PATH\n",
    "%env PATH={PATH}:/home/jupyter/.local/bin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f68df5dd-c456-4edd-8f58-71597f10c0ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retrieve and set PROJECT_ID and REGION environment variables.\n",
    "# TODO: Fill in the PROJECT_ID and REGION provided in the lab manual.\n",
    "PROJECT_ID = \"<PROJECT_ID>\"\n",
    "REGION = \"<REGION>\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a3912f9-6c12-439f-8613-cc60c286b3ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Create a globally unique Google Cloud Storage bucket for artifact storage.\n",
    "GCS_BUCKET = \"gs://<your bucket name>\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4931ae91-3ba1-437a-9c37-187a41a3d227",
   "metadata": {},
   "outputs": [],
   "source": [
    "!gcloud storage buckets create --location=$REGION $GCS_BUCKET"   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3ebbc2b-21ad-47f0-829f-9beba0deba9d",
   "metadata": {},
   "source": [
    "### Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bf558fc-d0fc-4452-8281-7d7cd0cffe50",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "import logging\n",
    "\n",
    "# TensorFlow model building libraries.\n",
    "import tensorflow as tf\n",
    "import tensorflow_text as text\n",
    "import tensorflow_hub as hub\n",
    "\n",
    "# Re-create the AdamW optimizer used in the original BERT paper.\n",
    "from official.nlp import optimization  \n",
    "\n",
    "# Libraries for data and plot model training metrics.\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Import the Vertex AI Python SDK.\n",
    "from google.cloud import aiplatform as vertexai\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d296167a-13b9-4895-be8b-b3b49fad5d47",
   "metadata": {},
   "source": [
    "### Initialize Vertex AI Python SDK"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34c178b0-0edb-4e4b-abb4-d3cc0bd676de",
   "metadata": {},
   "source": [
    "Initialize the Vertex AI Python SDK with your GCP Project, Region, and Google Cloud Storage Bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a43371e-2c64-4a76-8698-fa768043dbdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "vertexai.init(project=PROJECT_ID, location=REGION, staging_bucket=GCS_BUCKET)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2917411-811c-46dd-8eda-e8ef579c568d",
   "metadata": {},
   "source": [
    "## Build and train your model locally in a Vertex Notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8efc8cc5-ed5e-457a-b5f5-475bacca4611",
   "metadata": {},
   "source": [
    "Note: this lab adapts and extends the official [TensorFlow BERT text classification tutorial](https://www.tensorflow.org/text/tutorials/classify_text_with_bert#define_your_model) to utilize Vertex AI services. See the tutorial for additional coverage on fine-tuning BERT models using TensorFlow."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3338a818-18e5-4b0b-b37d-b387577a08ef",
   "metadata": {},
   "source": [
    "### Lab dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fdfb91d-6060-4d00-a1c3-299ee6027b76",
   "metadata": {},
   "source": [
    "In this lab, you will use the [Large Movie Review Dataset](https://ai.stanford.edu/~amaas/data/sentiment) that contains the text of 50,000 movie reviews from the Internet Movie Database. These are split into 25,000 reviews for training and 25,000 reviews for testing. The training and testing sets are balanced, meaning they contain an equal number of positive and negative reviews. Data ingestion and processing code has been provided for you below:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef491df4-c35f-4555-a6b6-96114c3d3c6e",
   "metadata": {},
   "source": [
    "### Import dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ee70d2c-c0e3-4c75-9bc6-b42dad6c7267",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_URL = \"https://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz\"\n",
    "LOCAL_DATA_DIR = \".\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c889f275-ce52-4108-9f7f-7cf824184f88",
   "metadata": {},
   "outputs": [],
   "source": [
    "def download_data(data_url, local_data_dir):\n",
    "    \"\"\"Download dataset.\n",
    "    Args:\n",
    "      data_url(str): Source data URL path.\n",
    "      local_data_dir(str): Local data download directory path.\n",
    "    Returns:\n",
    "      dataset_dir(str): Local unpacked data directory path.\n",
    "    \"\"\"\n",
    "    if not os.path.exists(local_data_dir):\n",
    "        os.makedirs(local_data_dir)\n",
    "    \n",
    "    dataset = tf.keras.utils.get_file(\n",
    "      fname=\"aclImdb_v1.tar.gz\",\n",
    "      origin=data_url,\n",
    "      untar=True,\n",
    "      cache_dir=local_data_dir,\n",
    "      cache_subdir=\"\")\n",
    "    \n",
    "    dataset_dir = os.path.join(os.path.dirname(dataset), \"aclImdb\")\n",
    "    \n",
    "    train_dir = os.path.join(dataset_dir, \"train\")\n",
    "    \n",
    "    # Remove unused folders to make it easier to load the data.\n",
    "    remove_dir = os.path.join(train_dir, \"unsup\")\n",
    "    shutil.rmtree(remove_dir)\n",
    "    \n",
    "    return dataset_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98f906a4-64a0-45ae-b376-757ef0f661fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET_DIR = download_data(data_url=DATA_URL, local_data_dir=LOCAL_DATA_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d95a61fa-cf55-470f-9837-c783c4bcccf8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a dictionary to iteratively add data pipeline and model training hyperparameters.\n",
    "HPARAMS = {\n",
    "    # Set a random sampling seed to prevent data leakage in data splits from files.\n",
    "    \"seed\": 42,\n",
    "    # Number of training and inference examples.\n",
    "    \"batch-size\": 32\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9aeea425-d288-44a7-9958-f2b1f48f9c34",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_datasets(dataset_dir, hparams):\n",
    "    \"\"\"Load pre-split tf.datasets.\n",
    "    Args:\n",
    "      hparams(dict): A dictionary containing model training arguments.\n",
    "    Returns:\n",
    "      raw_train_ds(tf.dataset): Train split dataset (20k examples).\n",
    "      raw_val_ds(tf.dataset): Validation split dataset (5k examples).\n",
    "      raw_test_ds(tf.dataset): Test split dataset (25k examples).\n",
    "    \"\"\"    \n",
    "\n",
    "    raw_train_ds = tf.keras.preprocessing.text_dataset_from_directory(\n",
    "        os.path.join(dataset_dir, 'train'),\n",
    "        batch_size=hparams['batch-size'],\n",
    "        validation_split=0.2,\n",
    "        subset='training',\n",
    "        seed=hparams['seed'])    \n",
    "\n",
    "    raw_val_ds = tf.keras.preprocessing.text_dataset_from_directory(\n",
    "        os.path.join(dataset_dir, 'train'),\n",
    "        batch_size=hparams['batch-size'],\n",
    "        validation_split=0.2,\n",
    "        subset='validation',\n",
    "        seed=hparams['seed'])\n",
    "\n",
    "    raw_test_ds = tf.keras.preprocessing.text_dataset_from_directory(\n",
    "        os.path.join(dataset_dir, 'test'),\n",
    "        batch_size=hparams['batch-size'])\n",
    "    \n",
    "    return raw_train_ds, raw_val_ds, raw_test_ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ff05aa4-d299-4c80-a29a-43c6bc3ac152",
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_train_ds, raw_val_ds, raw_test_ds = load_datasets(DATASET_DIR, HPARAMS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50ee40c0-9e37-483c-98f5-dcdb467a2bfc",
   "metadata": {},
   "outputs": [],
   "source": [
    "AUTOTUNE = tf.data.AUTOTUNE\n",
    "CLASS_NAMES = raw_train_ds.class_names\n",
    "\n",
    "train_ds = raw_train_ds.prefetch(buffer_size=AUTOTUNE)\n",
    "val_ds = raw_val_ds.prefetch(buffer_size=AUTOTUNE)\n",
    "test_ds = raw_test_ds.prefetch(buffer_size=AUTOTUNE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28f5734e-d97c-484d-9f52-f6fb4e153ef0",
   "metadata": {},
   "source": [
    "Let's print a few example reviews:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d794068-817c-4cb8-8e4c-c49860d0c92d",
   "metadata": {},
   "outputs": [],
   "source": [
    "for text_batch, label_batch in train_ds.take(1):\n",
    "  for i in range(3):\n",
    "    print(f'Review {i}: {text_batch.numpy()[i]}')\n",
    "    label = label_batch.numpy()[i]\n",
    "    print(f'Label : {label} ({CLASS_NAMES[label]})')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e3e6686-2fa1-453c-8259-8e5a87cba023",
   "metadata": {},
   "source": [
    "### Choose a pre-trained BERT model to fine-tune for higher accuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3502ad71-5747-4a11-9122-0b5c7b2049cd",
   "metadata": {},
   "source": [
    "[**Bidirectional Encoder Representations from Transformers (BERT)**](https://arxiv.org/abs/1810.04805v2) is a transformer-based text representation model pre-trained on massive datasets (3+ billion words) that can be fine-tuned for state-of-the art results on many natural language processing (NLP) tasks. Since release in 2018 by Google researchers, its has transformed the field of NLP research and come to form a core part of significant improvements to [Google Search](https://www.blog.google/products/search/search-language-understanding-bert). \n",
    "\n",
    "To meet your business requirements of achieving higher accuracy on a small dataset (20k training examples), you will use a technique called transfer learning to combine a pre-trained BERT encoder and classification layers to fine tune a new higher performing model for binary sentiment classification."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdd347cc-5f87-4b89-8833-6df850729ec3",
   "metadata": {},
   "source": [
    "For this lab, you will use a smaller BERT model that trades some accuracy for faster training times.\n",
    "\n",
    "The Small BERT models are instances of the original BERT architecture with a smaller number L of layers (i.e., residual blocks) combined with a smaller hidden size H and a matching smaller number A of attention heads, as published by\n",
    "\n",
    "Iulia Turc, Ming-Wei Chang, Kenton Lee, Kristina Toutanova: [\"Well-Read Students Learn Better: On the Importance of Pre-training Compact Models\"](https://arxiv.org/abs/1908.08962), 2019.\n",
    "\n",
    "They have the same general architecture but fewer and/or smaller Transformer blocks, which lets you explore tradeoffs between speed, size and quality.\n",
    "\n",
    "The following preprocessing and encoder models in the TensorFlow 2 SavedModel format use the implementation of BERT from the [TensorFlow Models Github repository](https://github.com/tensorflow/models/tree/master/official/nlp/bert) with the trained weights released by the authors of Small BERT."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22196658-1c30-49c7-8485-5d933fc8988e",
   "metadata": {},
   "outputs": [],
   "source": [
    "HPARAMS.update({\n",
    "    # TF Hub BERT modules.\n",
    "    \"tfhub-bert-preprocessor\": \"https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3\",\n",
    "    \"tfhub-bert-encoder\": \"https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-2_H-128_A-2/2\",\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc50b93a-df95-47d9-bc14-14502ae4eb54",
   "metadata": {},
   "source": [
    "Text inputs need to be transformed to numeric token ids and arranged in several Tensors before being input to BERT. TensorFlow Hub provides a matching preprocessing model for each of the BERT models discussed above, which implements this transformation using TF ops from the TF.text library. Since this text preprocessor is a TensorFlow model, It can be included in your model directly."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "780e50df-9d35-4116-a167-8353046bf6b9",
   "metadata": {},
   "source": [
    "For fine-tuning, you will use the same optimizer that BERT was originally trained with: the \"Adaptive Moments\" (Adam). This optimizer minimizes the prediction loss and does regularization by weight decay (not using moments), which is also known as [AdamW](https://arxiv.org/abs/1711.05101)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26396cb1-fc24-4e96-bef2-6fc8e2d500a6",
   "metadata": {},
   "source": [
    "For the learning rate `initial-learning-rate`, you will use the same schedule as BERT pre-training: linear decay of a notional initial learning rate, prefixed with a linear warm-up phase over the first 10% of training steps `n_warmup_steps`. In line with the BERT paper, the initial learning rate is smaller for fine-tuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b634139-a0d1-41e7-be23-c6e580a4f0e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "HPARAMS.update({\n",
    "    # Model training hyperparameters for fine tuning and regularization.\n",
    "    \"epochs\": 3,\n",
    "    \"initial-learning-rate\": 3e-5,\n",
    "    \"dropout\": 0.1 \n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e415aeb-5ab2-42ac-904a-ae0649d45a2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = HPARAMS['epochs']\n",
    "steps_per_epoch = tf.data.experimental.cardinality(train_ds).numpy()\n",
    "n_train_steps = steps_per_epoch * epochs\n",
    "n_warmup_steps = int(0.1 * n_train_steps)    \n",
    "\n",
    "OPTIMIZER = optimization.create_optimizer(init_lr=HPARAMS['initial-learning-rate'],\n",
    "                                          num_train_steps=n_train_steps,\n",
    "                                          num_warmup_steps=n_warmup_steps,\n",
    "                                          optimizer_type='adamw')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "005b4646-ce95-47c4-b1f7-886f59980386",
   "metadata": {},
   "source": [
    "### Build and compile a TensorFlow BERT sentiment classifier"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80920377-4596-4dbd-8eb7-8580327fdb24",
   "metadata": {},
   "source": [
    "Next, you will define and compile your model by assembling pre-built TF-Hub components and tf.keras layers.\n",
    "\n",
    "**Note:** For any help while defining the model, you can refer [**TensorFlow BERT text classification tutorial**](https://www.tensorflow.org/text/tutorials/classify_text_with_bert#define_your_model)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "289da96f-2aad-4c34-85ce-5916ea98778e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_text_classifier(hparams, optimizer):\n",
    "    \"\"\"Define and compile a TensorFlow BERT sentiment classifier.\n",
    "    Args:\n",
    "      hparams(dict): A dictionary containing model training arguments.\n",
    "    Returns:\n",
    "      model(tf.keras.Model): A compiled TensorFlow model.\n",
    "    \"\"\"\n",
    "    text_input = tf.keras.layers.Input(shape=(), dtype=tf.string, name='text')\n",
    "    # TODO: Add a hub.KerasLayer for BERT text preprocessing using the hparams dict. \n",
    "    # Name the layer 'preprocessing' and store in the variable 'preprocessor'.\n",
    "    \n",
    "    encoder_inputs = preprocessor(text_input)\n",
    "    # TODO: Add a trainable hub.KerasLayer for BERT text encoding using the hparams dict.\n",
    "    # Name the layer 'BERT_encoder' and store in the variable 'encoder'.\n",
    "\n",
    "    outputs = encoder(encoder_inputs)\n",
    "    # For the fine-tuning you are going to use the `pooled_output` array which represents \n",
    "    # each input sequence as a whole. The shape is [batch_size, H]. \n",
    "    # You can think of this as an embedding for the entire movie review.\n",
    "    classifier = outputs['pooled_output']\n",
    "    # Add dropout to prevent overfitting during model fine-tuning.\n",
    "    classifier = tf.keras.layers.Dropout(hparams['dropout'], name='dropout')(classifier)\n",
    "    classifier = tf.keras.layers.Dense(1, activation=None, name='classifier')(classifier)\n",
    "    model = tf.keras.Model(text_input, classifier, name='bert-sentiment-classifier')\n",
    "    \n",
    "    loss = tf.keras.losses.BinaryCrossentropy(from_logits=True)\n",
    "    metrics = tf.metrics.BinaryAccuracy()    \n",
    "    \n",
    "    model.compile(optimizer=optimizer,\n",
    "                  loss=loss,\n",
    "                  metrics=metrics)    \n",
    "    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "036b72cc-9e1c-49c9-8a90-7b09b6108f3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = build_text_classifier(HPARAMS, OPTIMIZER)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8198df2-c15a-4f79-a154-83c941aba5f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize your fine-tuned BERT sentiment classifier.\n",
    "tf.keras.utils.plot_model(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acfeb71d-4e19-4759-8f5c-293c8c7cee1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "TEST_REVIEW = ['this is such an amazing movie!']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1c2112f-9f04-470e-8636-38ea948cba9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "BERT_RAW_RESULT = model(tf.constant(TEST_REVIEW))\n",
    "print(BERT_RAW_RESULT)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d53dd9fa-d0a4-46ab-a123-3065b8fde7c8",
   "metadata": {},
   "source": [
    "### Train and evaluate your BERT sentiment classifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1f008fc-f696-4b71-9011-2a897d268795",
   "metadata": {},
   "outputs": [],
   "source": [
    "HPARAMS.update({\n",
    "    # TODO: Save your BERT sentiment classifier locally in the form of <key>:<path to save the model>. \n",
    "    # Hint: You can use the key as 'model-dir' and save it to './bert-sentiment-classifier-local'.\n",
    "    \n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9efd08f3-988f-408a-b694-af4702de85ba",
   "metadata": {},
   "source": [
    "**Note:** training your model locally will take about 10-15 minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24660956-d60a-4c25-a654-7b192a01a2d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_evaluate(hparams):\n",
    "    \"\"\"Train and evaluate TensorFlow BERT sentiment classifier.\n",
    "    Args:\n",
    "      hparams(dict): A dictionary containing model training arguments.\n",
    "    Returns:\n",
    "      history(tf.keras.callbacks.History): Keras callback that records training event history.\n",
    "    \"\"\"\n",
    "    # dataset_dir = download_data(data_url, local_data_dir)\n",
    "    raw_train_ds, raw_val_ds, raw_test_ds = load_datasets(DATASET_DIR, hparams)\n",
    "    \n",
    "    train_ds = raw_train_ds.cache().prefetch(buffer_size=AUTOTUNE)\n",
    "    val_ds = raw_val_ds.cache().prefetch(buffer_size=AUTOTUNE)\n",
    "    test_ds = raw_test_ds.cache().prefetch(buffer_size=AUTOTUNE)     \n",
    "    \n",
    "    epochs = hparams['epochs']\n",
    "    steps_per_epoch = tf.data.experimental.cardinality(train_ds).numpy()\n",
    "    n_train_steps = steps_per_epoch * epochs\n",
    "    n_warmup_steps = int(0.1 * n_train_steps)    \n",
    "    \n",
    "    optimizer = optimization.create_optimizer(init_lr=hparams['initial-learning-rate'],\n",
    "                                              num_train_steps=n_train_steps,\n",
    "                                              num_warmup_steps=n_warmup_steps,\n",
    "                                              optimizer_type='adamw')    \n",
    "    \n",
    "    mirrored_strategy = tf.distribute.MirroredStrategy()\n",
    "    with mirrored_strategy.scope():\n",
    "        model = build_text_classifier(hparams=hparams, optimizer=optimizer)\n",
    "    \n",
    "    logging.info(model.summary())\n",
    "        \n",
    "    history = model.fit(x=train_ds,\n",
    "                        validation_data=val_ds,\n",
    "                        epochs=epochs)  \n",
    "    \n",
    "    logging.info(\"Test accuracy: %s\", model.evaluate(test_ds))\n",
    "\n",
    "    # Export Keras model in TensorFlow SavedModel format.\n",
    "    model.save(hparams['model-dir'])\n",
    "    \n",
    "    return history"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0549e700-bf6a-415a-bb35-01150d9535e5",
   "metadata": {},
   "source": [
    "Based on the `History` object returned by `model.fit()`. You can plot the training and validation loss for comparison, as well as the training and validation accuracy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f59cab23-fbf0-44d2-9ee0-e2dd83390fa5",
   "metadata": {},
   "outputs": [],
   "source": [
    "history = train_evaluate(HPARAMS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91980420-8451-4869-b189-2b3693131ff8",
   "metadata": {},
   "outputs": [],
   "source": [
    "history_dict = history.history\n",
    "print(history_dict.keys())\n",
    "\n",
    "acc = history_dict['binary_accuracy']\n",
    "val_acc = history_dict['val_binary_accuracy']\n",
    "loss = history_dict['loss']\n",
    "val_loss = history_dict['val_loss']\n",
    "\n",
    "epochs = range(1, len(acc) + 1)\n",
    "fig = plt.figure(figsize=(10, 6))\n",
    "fig.tight_layout()\n",
    "\n",
    "plt.subplot(2, 1, 1)\n",
    "# \"bo\" is for \"blue dot\"\n",
    "plt.plot(epochs, loss, 'r', label='Training loss')\n",
    "# b is for \"solid blue line\"\n",
    "plt.plot(epochs, val_loss, 'b', label='Validation loss')\n",
    "plt.title('Training and validation loss')\n",
    "# plt.xlabel('Epochs')\n",
    "plt.ylabel('Loss')\n",
    "plt.legend()\n",
    "\n",
    "plt.subplot(2, 1, 2)\n",
    "plt.plot(epochs, acc, 'r', label='Training acc')\n",
    "plt.plot(epochs, val_acc, 'b', label='Validation acc')\n",
    "plt.title('Training and validation accuracy')\n",
    "plt.xlabel('Epochs')\n",
    "plt.ylabel('Accuracy')\n",
    "plt.legend(loc='lower right');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1fcbcff-18d5-4448-b695-fd6bf3189c2f",
   "metadata": {},
   "source": [
    "In this plot, the red lines represent the training loss and accuracy, and the blue lines are the validation loss and accuracy. Based on the plots above, you should see model accuracy of around 78-80% which exceeds your business requirements target of greater than 75% accuracy."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4fc3fed-aa4c-40b2-8c44-19be21ba4689",
   "metadata": {},
   "source": [
    "## Containerize your model code"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3905338-288b-4565-9abb-9053d7559315",
   "metadata": {},
   "source": [
    "Now that you trained and evaluated your model locally in a Vertex Notebook as part of an experimentation workflow, your next step is to train and deploy your model on Google Cloud's Vertex AI platform."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbbb61ec-cb3c-43bd-9d75-9d61ff52848e",
   "metadata": {},
   "source": [
    "To train your BERT classifier on Google Cloud, you will you will package your Python training scripts and write a Dockerfile that contains instructions on your ML model code, dependencies, and execution instructions. You will build your custom container with Cloud Build, whose instructions are specified in `cloudbuild.yaml` and publish your container to your Artifact Registry. This workflow gives you the opportunity to use the same container to run as part of a portable and scalable [Vertex Pipelines](https://cloud.google.com/vertex-ai/docs/pipelines/introduction) workflow. \n",
    "\n",
    "\n",
    "You will walk through creating the following project structure for your ML mode code:\n",
    "```\n",
    "|--/bert-sentiment-classifier\n",
    "   |--/trainer\n",
    "      |--__init__.py\n",
    "      |--model.py\n",
    "      |--task.py\n",
    "   |--Dockerfile\n",
    "   |--cloudbuild.yaml\n",
    "   |--requirements.txt\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a033e3c3-9dad-49d8-b53c-fd48113a8f90",
   "metadata": {},
   "source": [
    "### 1. Write a `model.py` training script\n",
    "\n",
    "First, you will tidy up your local TensorFlow model training code from above into a training script."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0129184d-15ed-4ebe-bbdc-6c2687eb18cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_DIR = \"bert-sentiment-classifier\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2594afe7-b9e0-4957-9156-d2e595fde62f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile {MODEL_DIR}/trainer/model.py\n",
    "import os\n",
    "import shutil\n",
    "import logging\n",
    "\n",
    "import tensorflow as tf\n",
    "import tensorflow_text as text\n",
    "import tensorflow_hub as hub\n",
    "from official.nlp import optimization\n",
    "\n",
    "DATA_URL = 'https://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz'\n",
    "LOCAL_DATA_DIR = './tmp/data'\n",
    "AUTOTUNE = tf.data.AUTOTUNE\n",
    "\n",
    "\n",
    "def download_data(data_url, local_data_dir):\n",
    "    \"\"\"Download dataset.\n",
    "    Args:\n",
    "      data_url(str): Source data URL path.\n",
    "      local_data_dir(str): Local data download directory path.\n",
    "    Returns:\n",
    "      dataset_dir(str): Local unpacked data directory path.\n",
    "    \"\"\"\n",
    "    if not os.path.exists(local_data_dir):\n",
    "        os.makedirs(local_data_dir)\n",
    "    \n",
    "    dataset = tf.keras.utils.get_file(\n",
    "      fname='aclImdb_v1.tar.gz',\n",
    "      origin=data_url,\n",
    "      untar=True,\n",
    "      cache_dir=local_data_dir,\n",
    "      cache_subdir=\"\")\n",
    "    \n",
    "    dataset_dir = os.path.join(os.path.dirname(dataset), 'aclImdb')\n",
    "    \n",
    "    train_dir = os.path.join(dataset_dir, 'train')\n",
    "    \n",
    "    # Remove unused folders to make it easier to load the data.\n",
    "    remove_dir = os.path.join(train_dir, 'unsup')\n",
    "    shutil.rmtree(remove_dir)\n",
    "    \n",
    "    return dataset_dir\n",
    "\n",
    "\n",
    "def load_datasets(dataset_dir, hparams):\n",
    "    \"\"\"Load pre-split tf.datasets.\n",
    "    Args:\n",
    "      hparams(dict): A dictionary containing model training arguments.\n",
    "    Returns:\n",
    "      raw_train_ds(tf.dataset): Train split dataset (20k examples).\n",
    "      raw_val_ds(tf.dataset): Validation split dataset (5k examples).\n",
    "      raw_test_ds(tf.dataset): Test split dataset (25k examples).\n",
    "    \"\"\"    \n",
    "\n",
    "    raw_train_ds = tf.keras.preprocessing.text_dataset_from_directory(\n",
    "        os.path.join(dataset_dir, 'train'),\n",
    "        batch_size=hparams['batch-size'],\n",
    "        validation_split=0.2,\n",
    "        subset='training',\n",
    "        seed=hparams['seed'])    \n",
    "\n",
    "    raw_val_ds = tf.keras.preprocessing.text_dataset_from_directory(\n",
    "        os.path.join(dataset_dir, 'train'),\n",
    "        batch_size=hparams['batch-size'],\n",
    "        validation_split=0.2,\n",
    "        subset='validation',\n",
    "        seed=hparams['seed'])\n",
    "\n",
    "    raw_test_ds = tf.keras.preprocessing.text_dataset_from_directory(\n",
    "        os.path.join(dataset_dir, 'test'),\n",
    "        batch_size=hparams['batch-size'])\n",
    "    \n",
    "    return raw_train_ds, raw_val_ds, raw_test_ds\n",
    "\n",
    "\n",
    "def build_text_classifier(hparams, optimizer):\n",
    "    \"\"\"Define and compile a TensorFlow BERT sentiment classifier.\n",
    "    Args:\n",
    "      hparams(dict): A dictionary containing model training arguments.\n",
    "    Returns:\n",
    "      model(tf.keras.Model): A compiled TensorFlow model.\n",
    "    \"\"\"\n",
    "    text_input = tf.keras.layers.Input(shape=(), dtype=tf.string, name='text')\n",
    "    # TODO: Add a hub.KerasLayer for BERT text preprocessing using the hparams dict. \n",
    "    # Name the layer 'preprocessing' and store in the variable 'preprocessor'.\n",
    "      \n",
    "    encoder_inputs = preprocessor(text_input)\n",
    "    # TODO: Add a trainable hub.KerasLayer for BERT text encoding using the hparams dict.\n",
    "    # Name the layer 'BERT_encoder' and store in the variable 'encoder'.\n",
    "      \n",
    "    outputs = encoder(encoder_inputs)\n",
    "    # For the fine-tuning you are going to use the `pooled_output` array which represents \n",
    "    # each input sequence as a whole. The shape is [batch_size, H]. \n",
    "    # You can think of this as an embedding for the entire movie review.\n",
    "    classifier = outputs['pooled_output']\n",
    "    # Add dropout to prevent overfitting during model fine-tuning.\n",
    "    classifier = tf.keras.layers.Dropout(hparams['dropout'], name='dropout')(classifier)\n",
    "    classifier = tf.keras.layers.Dense(1, activation=None, name='classifier')(classifier)\n",
    "    model = tf.keras.Model(text_input, classifier, name='bert-sentiment-classifier')\n",
    "    \n",
    "    loss = tf.keras.losses.BinaryCrossentropy(from_logits=True)\n",
    "    metrics = tf.metrics.BinaryAccuracy()    \n",
    "    \n",
    "    model.compile(optimizer=optimizer,\n",
    "                  loss=loss,\n",
    "                  metrics=metrics)    \n",
    "    \n",
    "    return model\n",
    "\n",
    "\n",
    "def train_evaluate(hparams):\n",
    "    \"\"\"Train and evaluate TensorFlow BERT sentiment classifier.\n",
    "    Args:\n",
    "      hparams(dict): A dictionary containing model training arguments.\n",
    "    Returns:\n",
    "      history(tf.keras.callbacks.History): Keras callback that records training event history.\n",
    "    \"\"\"\n",
    "    dataset_dir = download_data(data_url=DATA_URL, \n",
    "                                local_data_dir=LOCAL_DATA_DIR)\n",
    "    \n",
    "    raw_train_ds, raw_val_ds, raw_test_ds = load_datasets(dataset_dir=dataset_dir,\n",
    "                                                          hparams=hparams)\n",
    "    \n",
    "    train_ds = raw_train_ds.cache().prefetch(buffer_size=AUTOTUNE)\n",
    "    val_ds = raw_val_ds.cache().prefetch(buffer_size=AUTOTUNE)\n",
    "    test_ds = raw_test_ds.cache().prefetch(buffer_size=AUTOTUNE)     \n",
    "    \n",
    "    epochs = hparams['epochs']\n",
    "    steps_per_epoch = tf.data.experimental.cardinality(train_ds).numpy()\n",
    "    n_train_steps = steps_per_epoch * epochs\n",
    "    n_warmup_steps = int(0.1 * n_train_steps)    \n",
    "    \n",
    "    optimizer = optimization.create_optimizer(init_lr=hparams['initial-learning-rate'],\n",
    "                                              num_train_steps=n_train_steps,\n",
    "                                              num_warmup_steps=n_warmup_steps,\n",
    "                                              optimizer_type='adamw')    \n",
    "    \n",
    "    mirrored_strategy = tf.distribute.MirroredStrategy()\n",
    "    with mirrored_strategy.scope():\n",
    "        model = build_text_classifier(hparams=hparams, optimizer=optimizer)\n",
    "        logging.info(model.summary())\n",
    "        \n",
    "    history = model.fit(x=train_ds,\n",
    "                        validation_data=val_ds,\n",
    "                        epochs=epochs)  \n",
    "    \n",
    "    logging.info(\"Test accuracy: %s\", model.evaluate(test_ds))\n",
    "\n",
    "    # Export Keras model in TensorFlow SavedModel format.\n",
    "    model.save(hparams['model-dir'])\n",
    "    \n",
    "    return history"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e16b936-d93a-4411-a494-aacda93b05f4",
   "metadata": {},
   "source": [
    "### 2. Write a `task.py` file as an entrypoint to your custom model container"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17517e0b-a2ac-489a-bf03-357ace5d6577",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile {MODEL_DIR}/trainer/task.py\n",
    "\n",
    "import os\n",
    "import argparse\n",
    "\n",
    "from trainer import model\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    parser = argparse.ArgumentParser()\n",
    "    # Vertex custom container training args. These are set by Vertex AI during training but can also be overwritten.\n",
    "    parser.add_argument('--model-dir', dest='model-dir',\n",
    "                        default=os.environ['AIP_MODEL_DIR'], type=str, help='GCS URI for saving model artifacts.')\n",
    "\n",
    "    # Model training args.\n",
    "    parser.add_argument('--tfhub-bert-preprocessor', dest='tfhub-bert-preprocessor', \n",
    "                        default='https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3', type=str, help='TF-Hub URL.')\n",
    "    parser.add_argument('--tfhub-bert-encoder', dest='tfhub-bert-encoder', \n",
    "                        default='https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-2_H-128_A-2/2', type=str, help='TF-Hub URL.')    \n",
    "    parser.add_argument('--initial-learning-rate', dest='initial-learning-rate', default=3e-5, type=float, help='Learning rate for optimizer.')\n",
    "    parser.add_argument('--epochs', dest='epochs', default=3, type=int, help='Training iterations.')    \n",
    "    parser.add_argument('--batch-size', dest='batch-size', default=32, type=int, help='Number of examples during each training iteration.')    \n",
    "    parser.add_argument('--dropout', dest='dropout', default=0.1, type=float, help='Float percentage of DNN nodes [0,1] to drop for regularization.')    \n",
    "    parser.add_argument('--seed', dest='seed', default=42, type=int, help='Random number generator seed to prevent overlap between train and val sets.')\n",
    "    \n",
    "    args = parser.parse_args()\n",
    "    hparams = args.__dict__\n",
    "\n",
    "    model.train_evaluate(hparams)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d503a04a-5f15-4503-91e9-1acc4353fd08",
   "metadata": {},
   "source": [
    "### 3. Write a `Dockerfile` for your custom model container"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "253b0320-a3c2-4fbd-96f0-48ca3b49d485",
   "metadata": {},
   "source": [
    "Third, you will write a `Dockerfile` that contains instructions to package your model code in `bert-sentiment-classifier` as well as specifies your model code's dependencies needed for execution together in a Docker container."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b86ede10-6372-4320-89d3-264d4d1b1ca1",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile {MODEL_DIR}/Dockerfile\n",
    "# Specifies base image and tag.\n",
    "# https://cloud.google.com/vertex-ai/docs/training/pre-built-containers\n",
    "FROM us-docker.pkg.dev/vertex-ai/training/tf-cpu.2-11:latest\n",
    "\n",
    "# Sets the container working directory.\n",
    "WORKDIR /root\n",
    "\n",
    "# Copies the requirements.txt into the container to reduce network calls.\n",
    "COPY requirements.txt .\n",
    "\n",
    "# Installs additional packages.\n",
    "RUN pip3 install -U -r requirements.txt\n",
    "\n",
    "# b/203105209 Removes unneeded file from TF2.5 CPU image for python_module CustomJob training. \n",
    "# Will be removed on subsequent public Vertex images.\n",
    "RUN rm -rf /var/sitecustomize/sitecustomize.py\n",
    "\n",
    "# Copies the trainer code to the docker image.\n",
    "COPY . /trainer\n",
    "\n",
    "# Sets the container working directory.\n",
    "WORKDIR /trainer\n",
    "\n",
    "# Sets up the entry point to invoke the trainer.\n",
    "ENTRYPOINT [\"python\", \"-m\", \"trainer.task\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2974866-46f7-4f16-b6c0-9ea420ea6d73",
   "metadata": {},
   "source": [
    "### 4. Write a `requirements.txt` file to specify additional ML code dependencies"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7d62327-31c6-4ae5-819c-f59fb16e58c3",
   "metadata": {},
   "source": [
    "These are additional dependencies for your model code not included in the pre-built Vertex TensorFlow images such as TF-Hub, TensorFlow AdamW optimizer, and TensorFlow Text needed for importing and working with pre-trained TensorFlow BERT models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe7619e1-fff9-4a47-90a8-3ba9b55e74c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile {MODEL_DIR}/requirements.txt\n",
    "tf-models-official==2.15.0\n",
    "tensorflow-text==2.15.0\n",
    "tensorflow-hub==0.16.1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81292584-7a08-4c92-a9ba-e3dbc7005130",
   "metadata": {},
   "source": [
    "## Use Cloud Build to build and submit your model container to Google Cloud Artifact Registry"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d47400c-aa7c-4929-a584-f5067bd682eb",
   "metadata": {},
   "source": [
    "Next, you will use [Cloud Build](https://cloud.google.com/build) to build and upload your custom TensorFlow model container to [Google Cloud Artifact Registry](https://cloud.google.com/artifact-registry). \n",
    "\n",
    "Cloud Build brings reusability and automation to your ML experimentation by enabling you to reliably build, test, and deploy your ML model code as part of a CI/CD workflow. Artifact Registry provides a centralized repository for you to store, manage, and secure your ML container images. This will allow you to securely share your ML work with others and reproduce experiment results.\n",
    "\n",
    "**Note**: the initial build and submit step will take about 16 minutes but Cloud Build is able to take advantage of caching for faster subsequent builds. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "860c0d02-200f-4cc3-bdfd-ba96d233ecc4",
   "metadata": {},
   "source": [
    "### 1. Create Artifact Registry for custom container images \n",
    "\n",
    "**NOTE:** For any help to create the Artifact Registry, you can refer this [Documentation](https://cloud.google.com/sdk/gcloud/reference/artifacts/repositories/create)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9918475-f6dc-4fa3-8249-47976e68f529",
   "metadata": {},
   "outputs": [],
   "source": [
    "ARTIFACT_REGISTRY=\"bert-sentiment-classifier\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93d9566d-c6c4-48c8-b4f8-ad239e5ae349",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: create a Docker Artifact Registry using the gcloud CLI. Note the required 'repository-format', 'location' and 'description' flags while creating the Artifact Registry.\n",
    "# Documentation link: https://cloud.google.com/sdk/gcloud/reference/artifacts/repositories/create\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e900832e-de90-4ba1-ba7d-7973a1de9cc1",
   "metadata": {},
   "source": [
    "### 2. Create `cloudbuild.yaml` instructions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b580619d-957c-409f-ac15-ccbc0bb79a57",
   "metadata": {},
   "outputs": [],
   "source": [
    "IMAGE_NAME=\"bert-sentiment-classifier\"\n",
    "IMAGE_TAG=\"latest\"\n",
    "IMAGE_URI=f\"{REGION}-docker.pkg.dev/{PROJECT_ID}/{ARTIFACT_REGISTRY}/{IMAGE_NAME}:{IMAGE_TAG}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24790970-988e-4694-b8cd-a2d9d500c11b",
   "metadata": {},
   "outputs": [],
   "source": [
    "cloudbuild_yaml = f\"\"\"steps:\n",
    "- name: 'gcr.io/cloud-builders/docker'\n",
    "  args: [ 'build', '-t', '{IMAGE_URI}', '.' ]\n",
    "images: \n",
    "- '{IMAGE_URI}'\"\"\"\n",
    "\n",
    "with open(f\"{MODEL_DIR}/cloudbuild.yaml\", \"w\") as fp:\n",
    "    fp.write(cloudbuild_yaml)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c14b9c6-c120-482d-b5ab-c6a4c53bb205",
   "metadata": {},
   "source": [
    "### 3. Build and submit your container image to Artifact Registry using Cloud Build"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31e154b1-e584-496c-97b4-14795f80928b",
   "metadata": {},
   "source": [
    "**Note:** your custom model container will take about 16 minutes initially to build and submit to your Artifact Registry. Artifact Registry is able to take advantage of caching so subsequent builds take about 10 minutes. For any help to submit a build, you can refer this [**documentation**](https://cloud.google.com/sdk/gcloud/reference/builds/submit)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ebf0093-d66e-49c1-8a55-047020735e60",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: use Cloud Build to build and submit your custom model container to your Artifact Registry.\n",
    "# Documentation link: https://cloud.google.com/sdk/gcloud/reference/builds/submit\n",
    "# Hint: make sure the config flag is pointed at `{MODEL_DIR}/cloudbuild.yaml` defined above and you include your model directory as {MODEL_DIR}. Also, add a timeout flag.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cee35ac-ab83-472d-ab18-f622f3e3bc31",
   "metadata": {},
   "source": [
    "## Define a pipeline using the KFP V2 SDK"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5556979-3583-44fd-86df-d30fab8d9464",
   "metadata": {},
   "source": [
    "To address your business requirements and get your higher performing model into production to deliver value faster, you will define a pipeline using the [**Kubeflow Pipelines (KFP) V2 SDK**](https://www.kubeflow.org/docs/components/pipelines/sdk/v2/v2-compatibility) to orchestrate the training and deployment of your model on [**Vertex Pipelines**](https://cloud.google.com/vertex-ai/docs/pipelines) below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aef0e36b-3cb8-4660-bbb1-a5dcca49aebe",
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime\n",
    "# google_cloud_pipeline_components includes pre-built KFP components for interfacing with Vertex AI services.\n",
    "from google_cloud_pipeline_components import aiplatform as gcc_aip\n",
    "from kfp.v2 import dsl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c95f7070-d6e5-47ab-a860-d6a7e7892164",
   "metadata": {},
   "outputs": [],
   "source": [
    "TIMESTAMP=datetime.datetime.now().strftime('%Y%m%d%H%M%S')\n",
    "DISPLAY_NAME = \"bert-sentiment-{}\".format(TIMESTAMP)\n",
    "GCS_BASE_OUTPUT_DIR= f\"{GCS_BUCKET}/{MODEL_DIR}-{TIMESTAMP}\"\n",
    "\n",
    "USER = \"\"  # TODO: change this to your name.\n",
    "PIPELINE_ROOT = \"{}/pipeline_root/{}\".format(GCS_BUCKET, USER)\n",
    "\n",
    "print(f\"Model display name: {DISPLAY_NAME}\")\n",
    "print(f\"GCS dir for model training artifacts: {GCS_BASE_OUTPUT_DIR}\")\n",
    "print(f\"GCS dir for pipeline artifacts: {PIPELINE_ROOT}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4aecd42-c969-4ce0-a49a-9150c45a91e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pre-built Vertex model serving container for deployment.\n",
    "# https://cloud.google.com/vertex-ai/docs/predictions/pre-built-containers\n",
    "SERVING_IMAGE_URI = \"us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-11:latest\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f5e8dbc-04b3-4193-87f5-984d2b98a2d0",
   "metadata": {},
   "source": [
    "The pipeline consists of three components:\n",
    "\n",
    "* `CustomContainerTrainingJobRunOp` [(documentation)](https://google-cloud-pipeline-components.readthedocs.io/en/google-cloud-pipeline-components-0.2.0/google_cloud_pipeline_components.aiplatform.html#google_cloud_pipeline_components.aiplatform.CustomContainerTrainingJobRunOp): trains your custom model container using Vertex Training. This is the same as configuring a Vertex Custom Container Training Job using the Vertex Python SDK you covered in the Vertex AI: Qwik Start lab.\n",
    "\n",
    "*  `EndpointCreateOp` [(documentation)](https://google-cloud-pipeline-components.readthedocs.io/en/google-cloud-pipeline-components-0.2.0/google_cloud_pipeline_components.aiplatform.html#google_cloud_pipeline_components.aiplatform.EndpointCreateOp): Creates a Google Cloud Vertex Endpoint resource that maps physical machine resources with your model to enable it to serve online predictions. Online predictions have low latency requirements; providing resources to the model in advance reduces latency. \n",
    "\n",
    "* `ModelDeployOp`[(documentation)](https://google-cloud-pipeline-components.readthedocs.io/en/google-cloud-pipeline-components-0.2.0/google_cloud_pipeline_components.aiplatform.html#google_cloud_pipeline_components.aiplatform.ModelDeployOp): deploys your model to a Vertex Prediction Endpoint for online predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2181f3d-10cd-49c8-8e2f-e5c314940321",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dsl.pipeline(name=\"bert-sentiment-classification\", pipeline_root=PIPELINE_ROOT)\n",
    "def pipeline(\n",
    "    project: str = PROJECT_ID,\n",
    "    location: str = REGION,\n",
    "    staging_bucket: str = GCS_BUCKET,\n",
    "    display_name: str = DISPLAY_NAME,    \n",
    "    container_uri: str = IMAGE_URI,\n",
    "    model_serving_container_image_uri: str = SERVING_IMAGE_URI,    \n",
    "    base_output_dir: str = GCS_BASE_OUTPUT_DIR,\n",
    "):\n",
    "    \n",
    "    #TODO: add and configure the pre-built KFP CustomContainerTrainingJobRunOp component using\n",
    "    # the remaining arguments of the pipeline constructor as defined in the starting of this cell. \n",
    "    # Hint: Refer to the component documentation link above if needed as well.\n",
    "    model_train_evaluate_op = gcc_aip.CustomContainerTrainingJobRunOp(\n",
    "        # Vertex AI Python SDK authentication parameters.        \n",
    "        project=project,\n",
    "        location=location,\n",
    "        staging_bucket=staging_bucket,\n",
    "        # WorkerPool arguments.\n",
    "        replica_count=1,\n",
    "        machine_type=\"e2-standard-4\",\n",
    "        # TODO: fill in the remaining arguments from the pipeline definition.\n",
    "\n",
    "    )    \n",
    "    \n",
    "    # Create a Vertex Endpoint resource in parallel with model training.\n",
    "    endpoint_create_op = gcc_aip.EndpointCreateOp(\n",
    "        # Vertex AI Python SDK authentication parameters.\n",
    "        project=project,\n",
    "        location=location,\n",
    "        display_name=display_name\n",
    "    \n",
    "    )   \n",
    "    \n",
    "    # Deploy your model to the created Endpoint resource for online predictions.\n",
    "    model_deploy_op = gcc_aip.ModelDeployOp(\n",
    "        # Link to model training component through output model artifact.\n",
    "        model=model_train_evaluate_op.outputs[\"model\"],\n",
    "        # Link to the created Endpoint.\n",
    "        endpoint=endpoint_create_op.outputs[\"endpoint\"],\n",
    "        # Define prediction request routing. {\"0\": 100} indicates 100% of traffic \n",
    "        # to the ID of the current model being deployed.\n",
    "        traffic_split={\"0\": 100},\n",
    "        # WorkerPool arguments.        \n",
    "        dedicated_resources_machine_type=\"e2-standard-4\",\n",
    "        dedicated_resources_min_replica_count=1,\n",
    "        dedicated_resources_max_replica_count=2\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "783114fd-731b-4bad-bbe2-7a858e621fca",
   "metadata": {},
   "source": [
    "## Compile the pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb28dac2-3721-4fe6-9e01-98745b0d1aba",
   "metadata": {},
   "outputs": [],
   "source": [
    "from kfp.v2 import compiler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77355b83-577b-4831-9862-91e08e974256",
   "metadata": {},
   "outputs": [],
   "source": [
    "compiler.Compiler().compile(\n",
    "    pipeline_func=pipeline, package_path=\"bert-sentiment-classification.json\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "793cda30-4046-4d29-abdd-501c243f5eee",
   "metadata": {},
   "source": [
    "## Run the pipeline on Vertex Pipelines"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35be420d-9d1d-4e8e-a08a-658fdfd60eb0",
   "metadata": {},
   "source": [
    "The `PipelineJob` is configured below and triggered through the `run()` method.\n",
    "\n",
    "**Note:** This pipeline run will take around **30-40** minutes to train and deploy your model. Follow along with the execution using the URL from the job output below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f276575d-c2ba-4d08-9a2a-b7583af27aee",
   "metadata": {},
   "outputs": [],
   "source": [
    "vertex_pipelines_job = vertexai.pipeline_jobs.PipelineJob(\n",
    "    display_name=\"bert-sentiment-classification\",\n",
    "    template_path=\"bert-sentiment-classification.json\",\n",
    "    parameter_values={\n",
    "        \"project\": PROJECT_ID,\n",
    "        \"location\": REGION,\n",
    "        \"staging_bucket\": GCS_BUCKET,\n",
    "        \"display_name\": DISPLAY_NAME,        \n",
    "        \"container_uri\": IMAGE_URI,\n",
    "        \"model_serving_container_image_uri\": SERVING_IMAGE_URI,        \n",
    "        \"base_output_dir\": GCS_BASE_OUTPUT_DIR},\n",
    "    enable_caching=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0ab35e9-207c-49ea-8a27-6e6cddce8541",
   "metadata": {},
   "outputs": [],
   "source": [
    "vertex_pipelines_job.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "319a821a-a3bd-45bf-a9ea-aa18687218f6",
   "metadata": {},
   "source": [
    "## Query deployed model on Vertex Endpoint for online predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bfd8366-d362-4537-ab30-21c21fce6846",
   "metadata": {},
   "source": [
    "Finally, you will retrieve the `Endpoint` deployed by the pipeline and use it to query your model for online predictions.\n",
    "\n",
    "Configure the `Endpoint()` function below with the following parameters:\n",
    "\n",
    "*  `endpoint_name`: A fully-qualified endpoint resource name or endpoint ID. Example: \"projects/123/locations/us-central1/endpoints/456\" or \"456\" when project and location are initialized or passed.\n",
    "*  `project_id`: GCP project.\n",
    "*  `location`: GCP region.\n",
    "\n",
    "Call `predict()` to return a prediction for a test review."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf80748b-8907-4ad6-8adb-d4c394752257",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retrieve your deployed Endpoint name from your pipeline.\n",
    "ENDPOINT_NAME = vertexai.Endpoint.list()[0].name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c67c989d-1026-4f57-8dac-dafad01145a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "#TODO: Generate online predictions using your Vertex Endpoint. \n",
    "#Hint: You need to add the following variables: endpoint_name, project, location, with their required values.\n",
    "\n",
    "endpoint = vertexai.Endpoint(\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97759f45-e060-44ce-87fc-4d34c4b8cadf",
   "metadata": {},
   "outputs": [],
   "source": [
    "#TODO: write a movie review to test your model e.g. \"The Dark Knight is the best Batman movie!\"\n",
    "test_review = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71c008ce-90ad-4709-a24f-36b414c779e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: use your Endpoint to return prediction for your 'test_review' using 'endpoint.predict()' method.\n",
    "prediction ="
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54751a3e-7b2a-4ab8-b642-6533df27de82",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(prediction)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4a4c68e-a937-44f8-b64a-cbe9e607cb90",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use a sigmoid function to compress your model output between 0 and 1. For binary classification, a threshold of 0.5 is typically applied\n",
    "# so if the output is >= 0.5 then the predicted sentiment is \"Positive\" and < 0.5 is a \"Negative\" prediction.\n",
    "print(tf.sigmoid(prediction.predictions[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5344e3eb-0a0e-4271-b815-e792d8c95b66",
   "metadata": {},
   "source": [
    "## Next steps"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80557132-f0cf-4f4d-be5d-2c58453bc6b6",
   "metadata": {},
   "source": [
    "Congratulations! You walked through a full experimentation, containerization, and MLOps workflow on Vertex AI. First, you built, trained, and evaluated a BERT sentiment classifier model in a Vertex Notebook. You then packaged your model code into a Docker container to train on Google Cloud's Vertex AI. Lastly, you defined and ran a Kubeflow Pipeline on Vertex Pipelines that trained and deployed your model container to a Vertex Endpoint that you queried for online predictions."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6570ed8-a1ae-41e0-8a0b-9b63ca972d85",
   "metadata": {},
   "source": [
    "## License"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41c2c9ee-e982-4b8b-91c3-02f313896c6c",
   "metadata": {},
   "source": [
    "Copyright 2024 Google LLC\n",
    "\n",
    "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "you may not use this file except in compliance with the License.\n",
    "You may obtain a copy of the License at\n",
    "\n",
    "    https://www.apache.org/licenses/LICENSE-2.0\n",
    "\n",
    "Unless required by applicable law or agreed to in writing, software\n",
    "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "See the License for the specific language governing permissions and\n",
    "limitations under the License."
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "python3",
   "name": "tf2-gpu.2-6.m87",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/tf2-gpu.2-6:m87"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
